#!/usr/bin/env python3
import subprocess
import logging
import json
import os
import math

from boto3.session import Session

MAX_SPLIT_NUM = 100
NAS_ROOT_DIR = '/mnt/flow'

class FFmpegError(Exception):
    def __init__(self, message, status):
        super().__init__(message, status)
        self.message = message
        self.status = status

def exec_ffmpeg_cmd(cmd_lst):
    try:
        subprocess.run(
            cmd_lst, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    except subprocess.CalledProcessError as e:
        err_log = f'''
        returncode: {e.returncode}
        cmd: {e.cmd}
        output: {e.output}
        stderr: {e.stderr}
        stdout: {e.stdout}
        '''
        logging.error(err_log)
        raise FFmpegError(e.output, e.returncode)

def parse_file_path(file_path) -> tuple:
    (fileDir, file_name) = os.path.split(file_path)
    (shortname, extension) = os.path.splitext(file_name)
    return shortname, extension

def get_video_duration(input_video):
    cmd = '{0} -i {1} -show_entries format=duration -v quiet -of csv="p=0"'.format(
        '/opt/ffmpeg/ffprobe', input_video)
    raw_result = subprocess.check_output(cmd, shell=True)
    result = raw_result.decode().replace('\n', '').strip()
    duration = float(result)
    return duration

def handler(event, context):
    evt = json.loads(event)
    video_key = evt['video_key']
    oss_bucket_name = evt['oss_bucket_name']
    segment_time_seconds = str(evt['segment_time_seconds'])
    dst_formats = evt['dst_formats']

    access_key = os.getenv('ZOS_ACCESS_KEY')
    secret_key = os.getenv('ZOS_SECRET_KEY')
    endpoint = os.getenv('ZOS_ENDPOINT')
    session = Session(access_key, secret_key)
    s3_client = session.client('s3', endpoint_url=endpoint)

    _, file_name = os.path.split(video_key)
    shortname, extension = os.path.splitext(file_name)
    proc_dir = os.path.join(NAS_ROOT_DIR, context.request_id)
    os.makedirs(proc_dir)

    input_path = os.path.join(proc_dir, file_name)
    s3_client.download_file(oss_bucket_name, video_key, input_path)
    
    video_duration = get_video_duration(input_path)
    segment_time_seconds = int(segment_time_seconds)
    split_num = math.ceil(video_duration/segment_time_seconds)
    # adjust segment_time_seconds
    if split_num > MAX_SPLIT_NUM:
        segment_time_seconds = int(math.ceil(video_duration/MAX_SPLIT_NUM)) + 1
    
    segment_time_seconds = str(segment_time_seconds)
    exec_ffmpeg_cmd(['/opt/ffmpeg/ffmpeg', '-i', input_path, '-c', 'copy', '-f', 'segment', '-segment_time',
                     segment_time_seconds, '-reset_timestamps', '1', proc_dir + '/split_' + shortname + '_piece_%02d' + extension])

    split_keys = []
    for filename in os.listdir(proc_dir):
        if filename.startswith('split_' + shortname):
            file_path = os.path.join(proc_dir, filename)
            split_keys.append(file_path)

    new_dst_formats = []
    for dst_format in dst_formats:
        new_split_keys = []
        for split_key in split_keys:
            new_split_keys.append({
               'filekey': split_key,
               'video_proc_dir': proc_dir,
               'format': dst_format,
               'video_key': video_key,
            })

        new_dst_formats.append(
        {
          'format': dst_format,
          'video_proc_dir': proc_dir,
          'split_keys': new_split_keys
        })

    return {
        'video_proc_dir': proc_dir,
        'dst_formats': new_dst_formats
    }
